#include "MixedMaterial.h"
#include "core/Scene.h"

using namespace OBR;

std::shared_ptr<CudaProgram> MixedMaterial::program;
CudaBuffer MixedMaterial::sbt_data_payload;

MixedMaterial::MixedMaterial() : Material("shaders/Material/MixedMaterial.cu")
{
    this->program_desc.callables.entryFunctionNameDC = "__direct_callable__mixed_material";
    this->program_desc.kind = OPTIX_PROGRAM_GROUP_KIND_CALLABLES;
    this->program_desc.flags = OPTIX_GEOMETRY_FLAG_NONE;
}

MixedMaterial::~MixedMaterial()
{
}

void MixedMaterial::uploadVirtualProgramPayload(
    std::vector<std::shared_ptr<VirtualProgramEntry>> instances)
{
    // collect data
    MixedMaterialData *data = new MixedMaterialData[instances.size()];
    for (size_t i = 0; i < instances.size(); i++)
    {
        auto mat = std::dynamic_pointer_cast<MixedMaterial>(instances[i]);
        data[i].num_mat_mixed = mat->mat_list.size();
        float sum_weights = 0.0f;
        for (size_t idx = 0; idx < mat->mat_list.size() && idx < 6; idx++)
        {
            sum_weights += mat->weights[idx];
            data[i].mat_ids[idx] = mat->mat_list[idx]->program_id;
            data[i].mat_instance_ids[idx] = mat->mat_list[idx]->program_instance_id;
            data[i].weights[idx] = sum_weights; // this records the prefix sum of weights
        }
        // normalize weights
        for (size_t idx = 0; idx < mat->mat_list.size() && idx < 6; idx++)
        {
            data[i].weights[idx] /= sum_weights;
        }
        data[i].scale = sum_weights;
    }

    // set payload
    size_t payload_size = sizeof(MixedMaterialData) * instances.size();
    this->sbt_data_payload.resize(payload_size);
    this->sbt_data_payload.upload(data, payload_size);

    delete[] data;
}

void MixedMaterial::_buildFromConfig(const Json::Value &config)
{
    if (config["type"].asString().compare("mixed") != 0)
    {
        std::cerr << "wrong config value is sent to MixedMaterial::_buildFromConfig" << std::endl;
        std::cerr << "got:" << config << std::endl;
        throw std::runtime_error("invalid parameter");
    }
    for (Json::ArrayIndex i = 0; i < config["submaterials"].size(); i++)
    {
        Json::Value mat_conf = config["submaterials"][i];
        this->mat_ids.push_back(mat_conf["id"].asUInt());
        this->weights.push_back(mat_conf["weight"].asFloat());
    }
}

void MixedMaterial::postBuild(Scene &sc)
{
    for (size_t i = 0; i < this->mat_ids.size(); i++)
    {
        std::shared_ptr<VirtualProgramEntry> refmat = sc.materials[this->mat_ids[i]];
        this->mat_list.push_back(std::dynamic_pointer_cast<Material>(refmat));
    }
}

CUdeviceptr MixedMaterial::getPayloadDevicePtr()
{
    return MixedMaterial::sbt_data_payload.d_pointer();
}

void MixedMaterial::setProgram(std::shared_ptr<CudaProgram> prog)
{
    MixedMaterial::program = prog;
}

std::shared_ptr<CudaProgram> MixedMaterial::getProgram()
{
    return MixedMaterial::program;
}

OptixProgramGroupOptions &MixedMaterial::getProgramGroupOptions()
{
    return this->program_opt;
}

OptixProgramGroupDesc &MixedMaterial::getProgramGroupDesc()
{
    this->program_desc.callables.moduleDC = this->cuda_module->optix_module;
    return this->program_desc;
}